package main

import (
	"crypto/ecdsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/asn1"
	"encoding/hex"
	"encoding/pem"
	"errors"
	"fmt"
	"io/ioutil"
	"math/big"
	"strings"
)

func loadPublicKey(publicKey []byte) (*ecdsa.PublicKey, error) {
	// decode the key, assuming it's in PEM format
	block, _ := pem.Decode(publicKey)
	if block == nil {
		return nil, errors.New("Failed to decode PEM public key")
	}
	pub, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return nil, errors.New("Failed to parse ECDSA public key")
	}
	switch pub := pub.(type) {
	case *ecdsa.PublicKey:
		return pub, nil
	}
	return nil, errors.New("Unsupported public key type")
}

func loadPrivateKey(privKey []byte) (*ecdsa.PrivateKey, error) {
	paramEnd := strings.Index(string(privKey), "-----BEGIN EC PRIVATE KEY-----")
	priDstr := string(privKey)[paramEnd:]
	// decode the key, assuming it's in PEM format
	block, _ := pem.Decode([]byte(priDstr))
	if block == nil {
		return nil, errors.New("Failed to decode PEM private key")
	}
	pri, err := x509.ParseECPrivateKey(block.Bytes)
	if err != nil {
		return nil, errors.New("Failed to parse ECDSA private key" + err.Error())
	}
	return pri, nil
}

// Represents the two mathematical components of an ECDSA signature once
// decomposed.
type ECDSASignature struct {
	R, S *big.Int
}

// Encapsulates the overall message we're trying to decode and validate.
type Envelope struct {
	RawMessage     []byte `json:"message"`
	SignatureBytes []byte `json:"signature"`
}

// Helper function to compute the SHA256 hash of the given string of bytes.
func hash(b []byte) []byte {
	h := sha256.New()
	// hash the body bytes
	h.Write(b)
	// compute the SHA256 hash
	return h.Sum(nil)
}

// Attempts to create a new envelope structure from the given JSON string.
func NewEnvelopeFromMsgAndSign(msg, signData []byte) (*Envelope, error) {
	var e Envelope
	e.RawMessage = msg
	e.SignatureBytes = signData
	// now attempt to unmarshal the message body itself from the raw message
	return &e, nil
}

// The central validation routine that validates this message against the given
// public key. On success, returns nil, on failure returns a relevant error.
func (e *Envelope) Validate(publicKey *ecdsa.PublicKey) error {
	// unmarshal the R and S components of the ASN.1-encoded signature into our
	// signature data structure
	sig := &ECDSASignature{}
	_, err := asn1.Unmarshal(e.SignatureBytes, sig)
	if err != nil {
		fmt.Println(err)
		return err
	}
	// compute the SHA256 hash of our message
	h := hash(e.RawMessage)
	fmt.Println("sig.R", hex.EncodeToString(sig.R.Bytes()))
	fmt.Println("sig.S", hex.EncodeToString(sig.S.Bytes()))
	fmt.Println("sha256", hex.EncodeToString(h[:]))
	// validate the signature!
	valid := ecdsa.Verify(
		publicKey,
		h,
		sig.R,
		sig.S,
	)
	if !valid {
		return errors.New("Signature validation failed")
	}
	// signature is valid
	return nil
}

func main() {
	pub, err := ioutil.ReadFile("ECCkey.pem")
	if err != nil {
		fmt.Println(err)
		return
	}
	// extract the public key from the test key string
	publicKey, err := loadPublicKey(pub)
	if err != nil {
		fmt.Println("Failed to parse test public key:", err)
		return
	}
	publicKey.X.Bytes()
	sec := fmt.Sprintf("04%s%s", hex.EncodeToString(publicKey.X.Bytes()), hex.EncodeToString(publicKey.Y.Bytes()))
	b, err := hex.DecodeString(sec)
	if err != nil {
		fmt.Println(err)
		return
	}
	fmt.Println("pubkey:", sec)
	h := sha256.Sum256(b)
	fmt.Println("pubkey sha256:", hex.EncodeToString(h[:]))
	err = ioutil.WriteFile("./pub.pem", b, 0666)
	if err != nil {
		fmt.Println(err)
		return
	}
}
